package com.zxjbyte.yiyi.extension.socket.manager;

import com.zxjbyte.yiyi.framework.common.cache.caffeine.impl.CaffeineCache;
import com.zxjbyte.yiyi.framework.common.constant.HeaderMsg;
import com.zxjbyte.yiyi.framework.common.domain.message.StreamMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * WebSocket会话缓存管理器
 *
 * @Author zhangxingjia
 * @Date 2024/7/15 14:36
 * @Version: 1.0
 */
@Slf4j
public class WebSocketSessionManager {

    // 在线用户数统计
    private static final AtomicInteger onlineCount = new AtomicInteger(0);
    // 用户会话缓存
    private static final CaffeineCache<String, WebSocketSession> sessionCache = CaffeineCache.newCaffeineCache();

    // 添加会话
    public static void addSession(String userId, WebSocketSession session) {
        sessionCache.put(userId, session);
        onlineCount.getAndIncrement();
    }

    // 移除会话
    public static void removeByUserId(String userId) {
        sessionCache.remove(userId);
        onlineCount.getAndDecrement();
    }

    public static void removeSession(WebSocketSession session) {
        String userId = null;
        for(Map.Entry<String, WebSocketSession> entry : sessionCache.asMap().entrySet()) {
            if(entry.getValue().getId().equals(session.getId())) {
                userId = entry.getKey();
                break;
            }
        }
        if(userId != null) {
            removeByUserId(userId);
        }
    }

    // 获取会话
    public static WebSocketSession getSession(String userId) {
        return sessionCache.get(userId);
    }

    // 广播消息
    public static void broadcastMessage(String message) {
        broadcastMessage(StreamMessage.of(HeaderMsg.BROADCAST_MESSAGE, message));
    }

    public static void broadcastMessage(StreamMessage message) {
        if(message != null){
            sessionCache.asMap().values().forEach(session -> {
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(message.toJsonString()));
                    } catch (IOException e) {
                        log.error(e.getMessage(), e);
                    }
                }
            });
        }
    }

    /**
     * 指定用户发送消息
     *
     * @param userId
     * @param message
     * @return
     */
    public static boolean sendMessage(String userId, StreamMessage message) {
        WebSocketSession session = sessionCache.get(userId);
        if(session != null && session.isOpen() && message != null) {
            try {
                session.sendMessage(new TextMessage(message.toJsonString()));
            } catch (IOException e) {
                log.error(e.getMessage(), e);
                return false;
            }
        }
        return true;
    }

    public static boolean sendMessage(String userId, String message) {
        return sendMessage(userId, StreamMessage.of(HeaderMsg.SEND_MESSAGE, message));
    }

    // 获取在线用户数
    public static int getOnlineCount() {
        return onlineCount.get();
    }
}
